package us.codecraft.webmagic.scheduler;

import java.net.UnknownHostException;

import us.codecraft.webmagic.Request;
import us.codecraft.webmagic.Task;
import us.codecraft.webmagic.utils.MongoUtil;

import com.mongodb.BasicDBObject;
import com.mongodb.DB;
import com.mongodb.DBCollection;
import com.mongodb.DBCursor;
import com.mongodb.DBObject;
import com.mongodb.Mongo;
import com.mongodb.MongoException;

public class MongoScheduler implements Scheduler {
	Mongo mc;
	DB db;
	DBCollection coll;
	
	public MongoScheduler (String MongoHost,String MongoDB,String Collection) throws UnknownHostException, MongoException {
		mc = new Mongo(MongoHost);
		db = mc.getDB(MongoDB);
		coll = db.getCollection(Collection);
	}

	@Override
	public void push(Request request, Task task) {
		// TODO Auto-generated method stub
		try {
	        java.util.Date date= new java.util.Date();
	        MongoUtil mu = new MongoUtil(mc,db,coll);
	        if (!mu.ValueExist("url", request.getUrl())) {

	        	BasicDBObject bdb = new BasicDBObject();
		        bdb.append("time", String.valueOf(date.getTime()));
		        bdb.append("url", request.getUrl());
		        bdb.append("visited", "0");
		        
		        coll.insert(bdb);
	        }
		} catch (MongoException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (UnknownHostException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
	}

	@Override
	public synchronized Request poll(Task task) {
		BasicDBObject query = new BasicDBObject();
		String url = null;
		
		query.put("visited", "0");

		DBCursor cursor = coll.find(query);
		
		if (cursor.hasNext()) {
			DBObject resultbdb = cursor.next();
			url = resultbdb.get("url").toString();
			
			BasicDBObject update = new BasicDBObject();
			update.append("$set", new BasicDBObject().append("visited", "1"));
		 
			BasicDBObject searchQuery = (BasicDBObject) resultbdb;
		 
			coll.update(searchQuery, update);
		}

		return new Request(url);
	}

}
